const std = @import("std");
const builtin = @import("builtin");
const csr = @import("riscv").csr;
const root = @import("kernel.zig");
const trap = @import("trap.zig");
const clint = @import("clint.zig");

const print = root.print;

extern fn switch_to(to: *Context) callconv(.c) void; // 返回类型不能是 noreturn

pub const Context = packed struct {
    // Layout must be kept in sync with switch_to
    // ignore x0
    ra: usize = 0,
    sp: usize,
    gp: usize = 0,
    tp: usize = 0,
    t0: usize = 0,
    t1: usize = 0,
    t2: usize = 0,
    s0: usize = 0,
    s1: usize = 0,
    a0: usize = 0,
    a1: usize = 0,
    a2: usize = 0,
    a3: usize = 0,
    a4: usize = 0,
    a5: usize = 0,
    a6: usize = 0,
    a7: usize = 0,
    s2: usize = 0,
    s3: usize = 0,
    s4: usize = 0,
    s5: usize = 0,
    s6: usize = 0,
    s7: usize = 0,
    s8: usize = 0,
    s9: usize = 0,
    s10: usize = 0,
    s11: usize = 0,
    t3: usize = 0,
    t4: usize = 0,
    t5: usize = 0,
    t6: usize = 0,
    // upon is trap frame

    // save the pc to run in next schedule cycle
    pc: usize, // offset: 31 * @sizeOf(usize)

    pub fn init(stack: []u8, func: *const fn () void) Context {
        const sp = @intFromPtr(stack.ptr) + stack.len;
        std.debug.assert(sp % 16 == 0);

        return .{
            .sp = sp,
            .pc = @intFromPtr(func),
        };
    }
};

/// current is used to point to the context of current task
var current: u1 = 0;
const MAX_TASKS = std.math.pow(usize, 2, @typeInfo(@TypeOf(current)).int.bits);

const STACK_SIZE = 16 * switch (builtin.mode) {
    .Debug => 200,
    .ReleaseSafe => 66,
    .ReleaseFast => 48,
    .ReleaseSmall => 32,
};

/// 栈太小的话，会覆盖 context 的空间，mepc 地址不对导致执行 mret 异常
var ctx_tasks: [MAX_TASKS]Context = undefined;

/// In the standard RISC-V calling convention, the stack pointer sp
/// is always 16-byte aligned.
var stacks: [MAX_TASKS][STACK_SIZE]u8 align(16) = undefined;

pub fn schedule() void { // 返回类型不能是 noreturn
    current +%= 1;
    switch_to(&ctx_tasks[current]);
}

pub inline fn yield() void {
    // trigger a machine-level software interrupt
    const hart = 0;
    clint.software.trigger(hart);
}

pub fn init() void {
    // 初始化 mscratch
    csr.mscratch.write(0);

    inline for (.{ user_task0, user_task1 }, 0..) |func, i| {
        ctx_tasks[i] = .init(&stacks[i], func);
    }
}

const UserData = struct {
    counter: usize,
    str: []const u8,
};

pub var person: UserData = .{ .counter = 0, .str = "Jack" };

pub fn callback(arg: *anyopaque) void {
    var data: *UserData = @ptrCast(@alignCast(arg));
    data.counter += 1;
    print("======> TIMEOUT: {s}: {}\n", .{ data.str, data.counter });
}

comptime {
    asm (
        \\.global gethid
        \\gethid:
        \\    li a7, 0
        \\    ecall
        \\    ret
    );
}
extern fn gethid(hartid: *usize) usize;

fn user_task0() void {
    print("Task 0: Created!\n", .{});
    //yield();

    var hartid: usize = 65;
    print("ptr: 0x{x}\n", .{@intFromPtr(&hartid)});

    _ = gethid(&hartid);
    print("hartid: " ++ "\x1b[32m" ++ "{}\n" ++ "\x1b[m", .{hartid});

    var count: u4 = 0;

    while (true) {
        count +%= 1;
        print("Task 0: Running... {}\n", .{count});
        if (count % 7 == 0) trap.tests(); // exception!!
        delay(1000);
    }
}

fn user_task1() void {
    print("Task 1: Created!\n", .{});

    while (true) {
        print("Task 1: Running...\n", .{});
        delay(1000);
    }
}

fn delay(x: usize) void {
    var i: usize = 50000 * x;
    const ptr: *volatile usize = &i;
    while (ptr.* > 0) {
        ptr.* -= 1;
    }
}
